
#!/usr/bin/env python
# -*- coding: utf-8 -*-
import backtrader as bt
import pandas as pd
import numpy as np
from datetime import datetime

class ValueMomentumStrategy(bt.Strategy):
    params = (
        ('value_rank_period', 252),  # 价值因子计算周期（日）
        ('momentum_period', 63),     # 动量因子计算周期（日）
        ('rebalance_freq', 21),      # 调仓频率（日）
        ('top_n', 30),               # 持仓股票数量
    )

    def __init__(self):
        # 价值因子：PB倒数（越低越好）
        self.value_factor = {}
        # 动量因子：过去63日收益率（越高越好）
        self.momentum_factor = {}
        self.counter = 0

    def prenext(self):
        self.next()

    def next(self):
        self.counter += 1
        if self.counter % self.p.rebalance_freq != 0:
            return

        # 计算所有股票因子
        stocks = self.datas
        for d in stocks:
            # 价值因子计算（使用PB倒数）
            pb = d.close[0] / d.book_value[0]
            self.value_factor[d] = 1 / pb if pb > 0 else -np.inf
            
            # 动量因子计算
            ret = (d.close[0] - d.close[-self.p.momentum_period]) / d.close[-self.p.momentum_period]
            self.momentum_factor[d] = ret

        # 因子标准化与合成
        value_rank = {k: v for k, v in sorted(self.value_factor.items(), key=lambda item: item[1], reverse=True)}
        mom_rank = {k: v for k, v in sorted(self.momentum_factor.items(), key=lambda item: item[1], reverse=True)}
        
        # 综合得分 = 价值排名 + 动量排名
        composite_score = {}
        for d in stocks:
            composite_score[d] = (
                list(value_rank.keys()).index(d) + 
                list(mom_rank.keys()).index(d)
            )

        # 选择得分最高的top_n只股票
        selected = sorted(composite_score.items(), key=lambda x: x[1])[:self.p.top_n]
        selected_stocks = [d for d, _ in selected]

        # 执行调仓
        for d in stocks:
            if d in selected_stocks:
                self.order_target_percent(d, target=1.0/self.p.top_n)
            else:
                self.order_target_percent(d, target=0)

if __name__ == '__main__':
    cerebro = bt.Cerebro()
    cerebro.addstrategy(ValueMomentumStrategy)

    # 示例：加载沪深300成分股数据（需替换为实际数据）
    stocks_data = []
    for ticker in ['600519.SS', '000858.SZ', '601318.SS']:  # 示例股票
        df = pd.read_csv(f'{ticker}.csv', parse_dates=['date'], index_col='date')
        df['book_value'] = df['book_value'].ffill()  # 填充账面价值
        data = bt.feeds.PandasData(dataname=df)
        stocks_data.append(data)

    for d in stocks_data:
        cerebro.adddata(d)

    cerebro.broker.setcash(1000000)
    cerebro.broker.setcommission(commission=0.001)  # 0.1%交易佣金
    cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name='sharpe')
    cerebro.addanalyzer(bt.analyzers.DrawDown, _name='drawdown')

    print('初始资金: %.2f' % cerebro.broker.getvalue())
    results = cerebro.run()
    print('最终资金: %.2f' % cerebro.broker.getvalue())
    print('夏普比率:', results[0].analyzers.sharpe.get_analysis()['sharperatio'])
    print('最大回撤:', results[0].analyzers.drawdown.get_analysis()['max']['drawdown'])
